{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载调度器与模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Diffusion管道是由可互换的调度器(schedulers)和模型(models)组成的集合，可通过混合搭配来定制特定用例的流程。调度器封装了整个去噪过程（如去噪步数和寻找去噪样本的算法），其本身不包含可训练参数，因此内存占用极低。模型则主要负责从含噪输入到较纯净样本的前向传播过程。\n",
    "\n",
    "本指南将展示如何加载调度器和模型来自定义流程。我们将全程使用[stable-diffusion-v1-5/stable-diffusion-v1-5](https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5)检查点，首先加载基础管道："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import DiffusionPipeline\n",
    "\n",
    "pipeline = DiffusionPipeline.from_pretrained(\n",
    "    \"stable-diffusion-v1-5/stable-diffusion-v1-5\", torch_dtype=torch.float16, use_safetensors=True\n",
    ").to(\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通过`pipeline.scheduler`属性可查看当前管道使用的调度器："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.scheduler\n",
    "PNDMScheduler {\n",
    "  \"_class_name\": \"PNDMScheduler\",\n",
    "  \"_diffusers_version\": \"0.21.4\",\n",
    "  \"beta_end\": 0.012,\n",
    "  \"beta_schedule\": \"scaled_linear\",\n",
    "  \"beta_start\": 0.00085,\n",
    "  \"clip_sample\": false,\n",
    "  \"num_train_timesteps\": 1000,\n",
    "  \"set_alpha_to_one\": false,\n",
    "  \"skip_prk_steps\": true,\n",
    "  \"steps_offset\": 1,\n",
    "  \"timestep_spacing\": \"leading\",\n",
    "  \"trained_betas\": null\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 加载调度器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "调度器通过配置文件定义，同一配置文件可被多种调度器共享。使用`SchedulerMixin.from_pretrained()`方法加载时，需指定`subfolder`参数以定位配置文件在仓库中的正确子目录。\n",
    "\n",
    "例如加载`DDIMScheduler`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import DDIMScheduler, DiffusionPipeline\n",
    "\n",
    "ddim = DDIMScheduler.from_pretrained(\"stable-diffusion-v1-5/stable-diffusion-v1-5\", subfolder=\"scheduler\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后将新调度器传入管道："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = DiffusionPipeline.from_pretrained(\n",
    "    \"stable-diffusion-v1-5/stable-diffusion-v1-5\", scheduler=ddim, torch_dtype=torch.float16, use_safetensors=True\n",
    ").to(\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 调度器对比"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "不同调度器各有优劣，难以定量评估哪个最适合您的流程。通常需要在去噪速度与质量之间权衡。我们建议尝试多种调度器以找到最佳方案。通过`pipeline.scheduler.compatibles`属性可查看兼容当前管道的所有调度器。\n",
    "\n",
    "下面我们使用相同提示词和随机种子，对比`LMSDiscreteScheduler`、`EulerDiscreteScheduler`、`EulerAncestralDiscreteScheduler`和`DPMSolverMultistepScheduler`的表现："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import DiffusionPipeline\n",
    "\n",
    "pipeline = DiffusionPipeline.from_pretrained(\n",
    "    \"stable-diffusion-v1-5/stable-diffusion-v1-5\", torch_dtype=torch.float16, use_safetensors=True\n",
    ").to(\"cuda\")\n",
    "\n",
    "prompt = \"A photograph of an astronaut riding a horse on Mars, high resolution, high definition.\"\n",
    "generator = torch.Generator(device=\"cuda\").manual_seed(8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用`from_config()`方法加载不同调度器的配置来切换管道调度器：\n",
    "\n",
    "<hfoptions id=\"schedulers\">\n",
    "<hfoption id=\"LMSDiscreteScheduler\">\n",
    "\n",
    "`LMSDiscreteScheduler`通常能生成比默认调度器更高质量的图像。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import LMSDiscreteScheduler\n",
    "\n",
    "pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)\n",
    "image = pipeline(prompt, generator=generator).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "</hfoption>\n",
    "<hfoption id=\"EulerDiscreteScheduler\">\n",
    "\n",
    "`EulerDiscreteScheduler`仅需30步即可生成高质量图像。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import EulerDiscreteScheduler\n",
    "\n",
    "pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)\n",
    "image = pipeline(prompt, generator=generator).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "</hfoption>\n",
    "<hfoption id=\"EulerAncestralDiscreteScheduler\">\n",
    "\n",
    "`EulerAncestralDiscreteScheduler`同样可在30步内生成高质量图像。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import EulerAncestralDiscreteScheduler\n",
    "\n",
    "pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)\n",
    "image = pipeline(prompt, generator=generator).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "</hfoption>\n",
    "<hfoption id=\"DPMSolverMultistepScheduler\">\n",
    "\n",
    "`DPMSolverMultistepScheduler`在速度与质量间取得平衡，仅需20步即可生成优质图像。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import DPMSolverMultistepScheduler\n",
    "\n",
    "pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)\n",
    "image = pipeline(prompt, generator=generator).images[0]\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "</hfoption>\n",
    "</hfoptions>\n",
    "\n",
    "<div class=\"flex gap-4\">\n",
    "  <div>\n",
    "    <img class=\"rounded-xl\" src=\"https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/diffusers_docs/astronaut_lms.png\" />\n",
    "    <figcaption class=\"mt-2 text-center text-sm text-gray-500\">LMSDiscreteScheduler</figcaption>\n",
    "  </div>\n",
    "  <div>\n",
    "    <img class=\"rounded-xl\" src=\"https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/diffusers_docs/astronaut_euler_discrete.png\" />\n",
    "    <figcaption class=\"mt-2 text-center text-sm text-gray-500\">EulerDiscreteScheduler</figcaption>\n",
    "  </div>\n",
    "</div>\n",
    "<div class=\"flex gap-4\">\n",
    "  <div>\n",
    "    <img class=\"rounded-xl\" src=\"https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/diffusers_docs/astronaut_euler_ancestral.png\" />\n",
    "    <figcaption class=\"mt-2 text-center text-sm text-gray-500\">EulerAncestralDiscreteScheduler</figcaption>\n",
    "  </div>\n",
    "  <div>\n",
    "    <img class=\"rounded-xl\" src=\"https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/diffusers_docs/astronaut_dpm.png\" />\n",
    "    <figcaption class=\"mt-2 text-center text-sm text-gray-500\">DPMSolverMultistepScheduler</figcaption>\n",
    "  </div>\n",
    "</div>\n",
    "\n",
    "多数生成图像质量相近，实际选择需根据具体场景测试多种调度器进行比较。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Flax调度器"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对比Flax调度器时，需额外将调度器状态加载到模型参数中。例如将`FlaxStableDiffusionPipeline`的默认调度器切换为超高效的`FlaxDPMSolverMultistepScheduler`：\n",
    "\n",
    "> [!警告]\n",
    "> `FlaxLMSDiscreteScheduler`和`FlaxDDPMScheduler`目前暂不兼容`FlaxStableDiffusionPipeline`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import numpy as np\n",
    "from flax.jax_utils import replicate\n",
    "from flax.training.common_utils import shard\n",
    "from diffusers import FlaxStableDiffusionPipeline, FlaxDPMSolverMultistepScheduler\n",
    "\n",
    "scheduler, scheduler_state = FlaxDPMSolverMultistepScheduler.from_pretrained(\n",
    "    \"stable-diffusion-v1-5/stable-diffusion-v1-5\",\n",
    "    subfolder=\"scheduler\"\n",
    ")\n",
    "pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(\n",
    "    \"stable-diffusion-v1-5/stable-diffusion-v1-5\",\n",
    "    scheduler=scheduler,\n",
    "    variant=\"bf16\",\n",
    "    dtype=jax.numpy.bfloat16,\n",
    ")\n",
    "params[\"scheduler\"] = scheduler_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "利用Flax对TPU的兼容性实现并行图像生成。需为每个设备复制模型参数，并分配输入数据："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 每个并行设备生成1张图像（TPUv2-8/TPUv3-8支持8设备并行）\n",
    "prompt = \"一张宇航员在火星上骑马的高清照片，高分辨率，高画质。\"\n",
    "num_samples = jax.device_count()\n",
    "prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)\n",
    "\n",
    "prng_seed = jax.random.PRNGKey(0)\n",
    "num_inference_steps = 25\n",
    "\n",
    "# 分配输入和随机种子\n",
    "params = replicate(params)\n",
    "prng_seed = jax.random.split(prng_seed, jax.device_count())\n",
    "prompt_ids = shard(prompt_ids)\n",
    "\n",
    "images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images\n",
    "images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型加载"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通过`ModelMixin.from_pretrained()`方法加载模型，该方法会下载并缓存模型权重和配置的最新版本。若本地缓存已存在最新文件，则直接复用缓存而非重复下载。\n",
    "\n",
    "通过`subfolder`参数可从子目录加载模型。例如[stable-diffusion-v1-5/stable-diffusion-v1-5](https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5)的模型权重存储在[unet](https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5/tree/main/unet)子目录中："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import UNet2DConditionModel\n",
    "\n",
    "unet = UNet2DConditionModel.from_pretrained(\"stable-diffusion-v1-5/stable-diffusion-v1-5\", subfolder=\"unet\", use_safetensors=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "也可直接从[仓库](https://huggingface.co/google/ddpm-cifar10-32/tree/main)加载："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import UNet2DModel\n",
    "\n",
    "unet = UNet2DModel.from_pretrained(\"google/ddpm-cifar10-32\", use_safetensors=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "加载和保存模型变体时，需在`ModelMixin.from_pretrained()`和`ModelMixin.save_pretrained()`中指定`variant`参数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import UNet2DConditionModel\n",
    "\n",
    "unet = UNet2DConditionModel.from_pretrained(\n",
    "    \"stable-diffusion-v1-5/stable-diffusion-v1-5\", subfolder=\"unet\", variant=\"non_ema\", use_safetensors=True\n",
    ")\n",
    "unet.save_pretrained(\"./local-unet\", variant=\"non_ema\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用`from_pretrained()`的`torch_dtype`参数指定模型加载精度："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import AutoModel\n",
    "\n",
    "unet = AutoModel.from_pretrained(\n",
    "    \"stabilityai/stable-diffusion-xl-base-1.0\", subfolder=\"unet\", torch_dtype=torch.float16\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "也可使用[torch.Tensor.to](https://docs.pytorch.org/docs/stable/generated/torch.Tensor.to.html)方法即时转换精度，但会转换所有权重（不同于`torch_dtype`参数会保留`_keep_in_fp32_modules`中的层）。这对某些必须保持fp32精度的层尤为重要（参见[示例](https://github.com/huggingface/diffusers/blob/f864a9a352fa4a220d860bfdd1782e3e5af96382/src/diffusers/models/transformers/transformer_wan.py#L374)）。"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
